import re
import cv2
import numpy as np
import matplotlib.pyplot as plt
import pytesseract
from skimage import feature, exposure, transform, morphology
from scipy.signal import find_peaks


def detect_scale_bar(gray_img, scale_region_size=(113, 472), fallback_nm=10.0):
    """比例尺检测（带边界保护）"""
    h, w = gray_img.shape
    bar_h, bar_w = scale_region_size

    # 计算安全区域
    top = max(h - bar_h, 0)
    bottom = min(h, 1920)  # 防止 bottom < top
    if bottom <= top:  # 极端情况下 fallback
        bottom = top + 1  # 至少 1 像素

    left = 0
    right = min(bar_w, w)
    if right <= left:
        right = left + 1

    region = gray_img[top:bottom, left:right]

    # 如果 region 为空，直接提示并 fallback
    if region.size == 0:
        print("⚠️ 比例尺区域截取失败，使用默认值")
        return fallback_nm, 1  # 随便给个像素宽度 1

    # 其余逻辑保持不变
    dark_mask = (region < 40).astype(np.uint8)
    contours, _ = cv2.findContours(dark_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    max_width = 0
    for cnt in contours:
        _, _, w_, _ = cv2.boundingRect(cnt)
        max_width = max(max_width, w_)

    config = '--psm 6'
    ocr_text = pytesseract.image_to_string(region, config=config)
    print(f"🔍 OCR识别内容: '{ocr_text.strip()}'")

    match = re.search(r"(\d+(\.\d+)?)\s*(nm)?", ocr_text.lower())
    if match:
        scale_nm_raw = float(match.group(1))
        allowed = [5, 10, 20, 50, 100, 200]
        scale_nm = min(allowed, key=lambda x: abs(x - scale_nm_raw))
        print(f"✅ 提取数值: {scale_nm_raw} → 匹配最接近合法值: {scale_nm} nm")
    else:
        print(f"⚠️ OCR 未识别到有效比例尺文字，使用默认值 {fallback_nm} nm")
        scale_nm = fallback_nm

    if max_width == 0:
        raise ValueError("比例尺检测失败：未找到明显黑条。")
    print(f"📏 比例尺条像素长度: {max_width} px")
    return scale_nm, max_width


def image_preprocessing(image_path):
    """图片预处理"""
    # 1. 读取图像并利用高斯模糊去噪
    img = cv2.imread(image_path, 0)
    if img is None:
        print(f"错误:无法加载图像，请检查路径:{image_path}")
        return
    blurred = cv2.GaussianBlur(img, (3, 3), 0)

    # 2. 自适应对比度增强
    p2, p98 = np.percentile(blurred, (4, 98))  # 4% 的像素值被视为异常暗，2% 的像素值被视为异常亮,排除
    enhanced = exposure.rescale_intensity(blurred, in_range=(p2, p98))  # 对非极端部分提高对比度

    # 3. 自适应边缘检测:检测图像明暗变化梯度，确定边缘（线条）位置
    sobel_x = cv2.Sobel(enhanced, cv2.CV_64F, 1, 0, ksize=3)  # 水平梯度检测
    sobel_y = cv2.Sobel(enhanced, cv2.CV_64F, 0, 1, ksize=3)  # 垂直梯度检测
    # 两方向梯度平方相加并开根号:计算梯度幅值，表示图像中每个像素点处亮度变化的强度，幅值大的位置对应图像中的边缘
    gradient_mag = np.sqrt(sobel_x ** 2 + sobel_y ** 2)

    grad_mean = np.mean(gradient_mag)  # 计算整个图像梯度幅值的平均值
    # 更具图像梯度幅值确定sigma值
    if grad_mean > 233:
        sigma = 5
    else:
        sigma = 3
    # 提高亮度
    edges = feature.canny(enhanced, sigma=sigma)
    # edges = morphology.closing(edges, footprint_rectangle((3, 3))) # 闭合孔隙，减少断线

    # 4. 显示中间结果:三张图进行对比
    plt.figure(figsize=(12, 8))
    plt.subplot(221), plt.imshow(img, cmap='gray'), plt.title('原图')
    plt.subplot(222), plt.imshow(enhanced, cmap='gray'), plt.title('对比度增强')
    plt.subplot(223), plt.imshow(edges, cmap='gray'), plt.title('边缘检测结果')
    plt.tight_layout()
    plt.show()

    return img, enhanced, edges


def manual_box_selection(img, edges):
    """手动框选条纹区域"""
    print("请在弹出窗口中框选包含条纹的矩形区域（尽量包含多条完整条纹），按回车确认:")
    roi = cv2.selectROI("框选条纹区域", img)  # 打开交互窗口，显示原图
    x, y, w, h = roi  # 获得左上角坐标和图像高度和宽度
    roi_img = img[y:y + h, x:x + w]  # 提取原图像区域
    roi_edges = edges[y:y + h, x:x + w]  # 提取edges图像区域
    cv2.destroyAllWindows()  # 关闭交互窗口
    return roi_img, roi_edges


def hough_transform(roi_edges, roi_img):
    """利用霍夫变换得出旋转角"""
    print("正在利用霍夫变换检测条纹方向...")
    roi_edges_uint8 = roi_edges.astype(np.uint8) * 255  # 图片类型转换，后续需要

    # 霍夫变换检测直线，返回张量shape为(ddd, 1, 4)，每(1, 1, 4)储存两个端点:[[[x1, y1, x2, y2]]]
    lines = cv2.HoughLinesP(
        roi_edges_uint8,
        rho=1,  # 距离分辨率（像素），表示霍夫空间中 ρ 的步长。
        theta=np.pi / 180,  # 角度分辨率（弧度），表示霍夫空间中 θ 的步长，这里设为 1 度。
        threshold=10,  # 累加器阈值，只有投票数超过该值的线段才会被保留。较低的阈值会增加检测到的直线数量。
        minLineLength=15,  # 最小线段长度（像素），短于该长度的线段会被过滤。
        maxLineGap=3  # 线段最大间隔（像素），若两条线段间距小于此值，会被视为同一条直线。
    )
    img_copied = roi_img.copy()
    # 绘制检测到的线段
    if lines is not None:
        if lines is not None:
            # 绘制霍夫变换检测到的直线
            plt.subplot(221)
            for line in lines:
                x1, y1, x2, y2 = line[0]  # 获得线段两端点的坐标
                # 绘制线段
                plt.plot((x1, x2), (-y1, -y2), 'k-', linewidth=2)  # 黑色线段
                # 绘制端点，使用不同颜色
                plt.plot(x1, -y1, 'ro', markersize=3)  # 红色起点
                plt.plot(x2, -y2, 'bo', markersize=3)  # 蓝色终点
            plt.title('霍夫变换检测到的直线')
            plt.grid(True)

            # 在原图上绘制检测到的直线
            plt.subplot(222)
            for line in lines:
                x1, y1, x2, y2 = line[0]  # 正确获取当前线段的坐标
                cv2.line(img_copied, (x1, y1), (x2, y2), (0, 255, 0), 1)
            plt.imshow(cv2.cvtColor(img_copied, cv2.COLOR_BGR2RGB))
            plt.title('霍夫检测结果')
            plt.tight_layout()  # 调整子图布局
            plt.show()

    if lines is not None and len(lines) > 3:
        angles_hough = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            dx = x2 - x1
            dy = -(y2 - y1)  # 坐标轴转换
            angle = np.degrees(np.arctan2(dy, dx))
            angle %= 180
            if angle > 90:
                angle -= 180
            elif angle < -90:
                angle += 180
            angles_hough.append(angle)
        if angles_hough:
            angles_rad = np.deg2rad(angles_hough)
            median_sin = np.median(np.sin(angles_rad))
            median_cos = np.median(np.cos(angles_rad))
            angle_hough = np.rad2deg(np.arctan2(median_sin, median_cos))
            rotation_angle = angle_hough
            print(f'霍夫变换顺时针旋转{rotation_angle:.2f}度')
            return rotation_angle, roi_edges_uint8
    return None


def remove_black_lines(image, max_black_length=10):
    """
    删除包含过长连续全黑像素的行

    参数:
    image (numpy.ndarray): 输入图像(灰度, uint8格式)
    max_black_length (int): 允许的最大连续全黑像素长度(默认10)

    返回:
    numpy.ndarray: 处理后的图像
    """
    # 创建一个掩码标记需要保留的行
    rows_to_keep = []

    # 遍历图像的每一行
    for i in range(image.shape[0]):
        row = image[i, :]

        # 计算当前行连续全黑像素的最大长度
        current_length = 0
        max_current_length = 0

        # 遍历当前行的每个像素
        for pixel in row:
            if pixel == 0:  # 全黑像素
                current_length += 1
            else:
                if current_length > max_current_length:
                    max_current_length = current_length
                current_length = 0

        # 检查行末尾的连续黑色
        if current_length > max_current_length:
            max_current_length = current_length

        # 如果最大连续黑色长度不超过阈值，则保留该行
        if max_current_length <= max_black_length:
            rows_to_keep.append(i)

    # 保留符合条件的行
    if rows_to_keep:
        return image[rows_to_keep, :]
    else:
        # 如果没有行符合条件，返回空图像
        return np.empty((0, image.shape[1]), dtype=image.dtype)


def rotate_image(roi_img, roi_edges_uint8, rotation_angle):
    """旋转图像，使直线垂直分布"""
    rotated_roi = transform.rotate(roi_img, -rotation_angle, mode='constant')
    rotated_edges = transform.rotate(roi_edges_uint8, -rotation_angle, mode='constant')
    rotated_roi = remove_black_lines(rotated_roi, max_black_length=10)
    return rotated_roi, rotated_edges


def image_projection(rotated_roi, scale_bar_actual, scale_bar_pixels):
    """水平投影"""
    rotated_roi = 255 - rotated_roi  # 图像取反
    projection = np.sum(rotated_roi, axis=1)

    # 优化峰值检测（减少误判）
    proj_mean = np.mean(projection)
    proj_std = np.std(projection)
    proj_max = np.max(projection)
    proj_min = np.min(projection)
    proj_range = proj_max - proj_min  # 投影对比度

    # 动态阈值（根据对比度调整）
    if proj_range > 50:
        threshold = proj_mean + 0.8 * proj_std
    elif proj_range > 30:
        threshold = proj_mean + 0.6 * proj_std
    else:
        threshold = proj_mean + 0.4 * proj_std

    # 增加峰值突出度过滤
    peaks, properties = find_peaks(
        projection,
        height=threshold,
        distance=10,
        prominence=10
    )

    if len(peaks) >= 2:
        pixel_spacings = np.diff(peaks)
        avg_pixel_spacing = np.mean(pixel_spacings)
        actual_spacing = avg_pixel_spacing * (scale_bar_actual / scale_bar_pixels)

        print(f"检测到 {len(peaks)} 个条纹，平均像素间隔:{avg_pixel_spacing:.2f} 像素")
        print(f"实际物理间隔:{actual_spacing:.2f} nm")

        plt.figure(figsize=(10, 6))
        plt.subplot(211), plt.imshow(rotated_roi, cmap='gray')
        plt.title(f'旋转后ROI - 间隔:{actual_spacing:.2f} nm')
        plt.subplot(212), plt.plot(projection)
        plt.axhline(y=threshold, color='r', linestyle='--', label='阈值')
        plt.plot(peaks, projection[peaks], 'ro', label='条纹峰值')
        plt.title('垂直投影结果')
        plt.legend()
        plt.tight_layout()
        plt.show()

        if len(peaks) < 5:
            print("警告:检测到的条纹数量较少，建议手动调整ROI或角度")


def main(image_path):
    img, enhanced, edges = image_preprocessing(image_path=image_path)
    scale_bar_actual, scale_bar_pixels = detect_scale_bar(img, fallback_nm=10.0)
    roi_img, roi_edges = manual_box_selection(img=img, edges=edges)
    rotation_angle, roi_edges_uint8 = hough_transform(roi_edges=roi_edges, roi_img=roi_img)
    rotated_roi, rotated_edges = rotate_image(roi_img=roi_img, roi_edges_uint8=roi_edges_uint8,
                                              rotation_angle=rotation_angle)
    image_projection(rotated_roi=rotated_roi, scale_bar_actual=scale_bar_actual, scale_bar_pixels=scale_bar_pixels)


if __name__ == "__main__":
    image_path = "..\\img\\new\\0006.tif"  # 图像路径
    plt.rc('font', family='Microsoft YaHei')
    main(image_path)
